{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yeadDkMiISin"
      },
      "source": [
        "# Gemini API: Streaming Quickstart"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6e5f3f6c8c11"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/quickstarts/Streaming.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df1767a3d1cc"
      },
      "source": [
        "This notebook demonstrates streaming in the Python SDK. By default, the Python SDK returns a response after the model completes the entire generation process. You can also stream the response as it is being generated, and the model will return chunks of the response as soon as they are generated.\n",
        "\n",
        "**Download this notebook and run it locally (not in Google Colab)**\n",
        "\n",
        "Streaming is not handled correctly in Google Colab yet. Currently all the stream chunks are returned together, not as they are generated. To see the correct behavior, download this notebook and run it locally using Jupyter, instead."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xuiLSV7amy3P"
      },
      "outputs": [],
      "source": [
        "!pip install -U -q google-generativeai # Install the Python SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "79EWm0DAmy-g"
      },
      "outputs": [],
      "source": [
        "import google.generativeai as genai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DkeZNMrw6kPD"
      },
      "source": [
        "You'll need an API key stored in an environment variable to run this notebook. See the the [Authentication quickstart](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) for an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t9O-OzeAKC_m"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "genai.configure(api_key=os.environ['GOOGLE_API_KEY'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BUoa5q0iUuE1"
      },
      "source": [
        "## Handle streaming responses\n",
        "\n",
        "To stream responses, use [`GenerativeModel.generate_content(..., stream=True)`](https://ai.google.dev/api/python/google/generativeai/GenerativeModel#generate_content).\n",
        "\n",
        "**Note**: This cell runs with a Google Colab runtime, but does not properly show streaming due to implementation details of Colab runtimes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nVWWGBsBok3m"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In the quaint little cottage nestled amidst sprawling meadows, there lived two adorable cats named\n",
            "________________________________________________________________________________\n",
            " Whiskers and Mittens. Whiskers, a sleek black feline with emerald-green eyes, was known for his playful antics and mischievous nature. Mittens,\n",
            "________________________________________________________________________________\n",
            " on the other hand, was a demure white cat with piercing blue eyes and a soft, fluffy coat. She often curled up in a ball, purring contentedly by the fireplace.\n",
            "\n",
            "One sunny afternoon, as Whiskers and Mittens basked in the warm sunlight streaming through the window, a playful thought crossed\n",
            "________________________________________________________________________________\n",
            " Whiskers' mischievous mind. He stealthily crept behind Mittens and gently nudged her with his velvety nose. Mittens, startled, jumped up and spun around, her tail twitching in amusement.\n",
            "\n",
            "\"Oh, hello, Whiskers,\" Mittens purred, her voice as sweet as honey. \"What a pleasant surprise.\"\n",
            "\n",
            "Whiskers feigned innocence and rubbed his head against her leg. \"My dear Mittens, I couldn't resist disturbing your sleepy slumber. How can you possibly resist such a charming feline?\"\n",
            "\n",
            "Mittens couldn't help but chuckle at Whiskers' antics. \"You sly cat\n",
            "________________________________________________________________________________\n",
            ",\" she said. \"You know you're the one who interrupted my nap.\"\n",
            "\n",
            "Unfazed, Whiskers continued to charm Mittens, weaving tales of his daring adventures and making her laugh with his silly jokes. Mittens, unable to resist his infectious enthusiasm, found herself drawn into his playful antics.\n",
            "\n",
            "As the sun began to set, casting a golden glow over the meadow, Whiskers and Mittens found themselves chasing each other through the fields. They leaped over fences, scampered through wildflowers, and chased butterflies with abandon. The playful laughter of the cats echoed through the air, mingling with the sweet scent of blooming roses.\n",
            "\n",
            "Finally, as darkness enveloped the meadow, Whiskers and Mittens made their way back to the cottage, exhausted but content. They curled up together on the soft rug in front of the fireplace, purring softly as they drifted off to sleep.\n",
            "\n",
            "In the cozy warmth of the cottage, surrounded by the love of her playful companion, Mittens realized that true happiness could be found in the simplest of moments. And so, as the night wore on, Whiskers and Mittens slept soundly, their bond unbreakable, their hearts filled with love and purrfect contentment.\n",
            "________________________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "model = genai.GenerativeModel('gemini-pro')\n",
        "response = model.generate_content(\"Write a cute story about cats.\", stream=True)\n",
        "for chunk in response:\n",
        "  print(chunk.text)\n",
        "  print(\"_\"*80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KswwVyHCKC_n"
      },
      "source": [
        "## Handle streaming responses asynchronously\n",
        "\n",
        "To stream responses asynchronously, use [`GenerativeModel.generate_content_async(..., stream=True)`](https://ai.google.dev/api/python/google/generativeai/GenerativeModel#generate_content_async).\n",
        "\n",
        "**Note**: These cells do NOT work with a Google Colab runtime, but do work in a local Jupyter notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n6sXnWrJoKoo"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The sun peeked through the window, casting golden rays upon the sleeping form of\n",
            "________________________________________________________________________________\n",
            " Mittens. Her soft, white fur gently rose and fell with each breath, her tiny paws twitching as she chased invisible mice in her dreams.\n",
            "\n",
            "Next\n",
            "________________________________________________________________________________\n",
            " to her, curled up in a cozy ball, was her mischievous companion, Whiskers. His velvety black fur glistened, and his emerald eyes sparkled with feline wisdom. With his tail curled around himself, he dozed contentedly, his gentle purring filling the air.\n",
            "\n",
            "As the sun climbed higher,\n",
            "________________________________________________________________________________\n",
            " it reached Mittens' eyelids, causing her to stir. She stretched lazily, her long, slender body unfurling like a delicate ribbon. A soft meow escaped her lips, and Whiskers opened one eye sleepily.\n",
            "\n",
            "\"Good morning, Mittens,\" he said, his voice a velvety whisper. \"Time for breakfast.\"\n",
            "\n",
            "Mittens purred and jumped down from the bed, her graceful movements a testament to her inherent agility. Whiskers followed suit, his nimble paws carrying him to the kitchen.\n",
            "\n",
            "As they approached their food bowls, Mittens couldn't help but notice something peculiar. \"Whiskers,\" she\n",
            "________________________________________________________________________________\n",
            " exclaimed, \"there's something... in my bowl.\"\n",
            "\n",
            "Whiskers peered into the bowl and gasped. There, nestled among the kibble, was a tiny, iridescent ball.\n",
            "\n",
            "\"Oh my stars!\" cried Whiskers. \"It's a bell!\"\n",
            "\n",
            "Mittens was overjoyed. \"That's perfect!\" she meowed. \"Now I can jingle my way around the house and make everyone smile.\"\n",
            "\n",
            "With newfound excitement, Mittens and Whiskers tucked into their breakfast. As they ate, the kitchen filled with a chorus of happy meows and the cheerful jingle of Mittens' bell.\n",
            "\n",
            "Throughout the day, the two cats reveled in their newfound accessory. Mittens proudly jingled her bell as she explored the house, while Whiskers couldn't resist batting it playfully with his paws. Together, they created a symphony of sounds that brought joy to everyone who heard them.\n",
            "\n",
            "As the sun began to set, Mittens and Whiskers curled up on the sofa, exhausted but content. The jingle of Mittens' bell gradually subsided, replaced by the gentle sound of their purring.\n",
            "\n",
            "And so, in the cozy confines of their home, the two feline companions drifted off to sleep, their hearts filled with love and the whimsical\n",
            "________________________________________________________________________________\n",
            " sound of a tiny bell.\n",
            "________________________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "async for chunk in await model.generate_content_async(\"Write a cute story about cats.\", stream=True):\n",
        "  print(chunk.text)\n",
        "  print(\"_\"*80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jpK3p1B4KC_o"
      },
      "source": [
        "Here's a simple example of two asynchronous functions running simultaneously."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IJ-8SjYwKC_o"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "==========not blocked!==========\n",
            "In the quaint little cottage nestled amidst a verdant meadow, resided two adorable cats\n",
            "________________________________________________________________________________\n",
            " named Mittens and Whiskers. Mittens, with her soft gray fur and emerald-green eyes, possessed a mischievous spark that made everyone smile. Whiskers\n",
            "________________________________________________________________________________\n",
            "==========not blocked!==========\n",
            ", on the other hand, was a dignified tuxedo cat with piercing blue eyes and an air of quiet confidence.\n",
            "\n",
            "One sunny afternoon, as the warm breeze carried the scent of blooming lilacs through the open window, Mittens and Whiskers found themselves lounging on the sun-drenched windowsill. As they basked\n",
            "________________________________________________________________________________\n",
            "==========not blocked!==========\n",
            " in the golden rays, their gaze fell upon a playful squirrel scampering along the branch of a nearby tree.\n",
            "\n",
            "Instantly, Mittens' hunter instincts ignited. She stealthily slid off the windowsill and approached the tree, her green eyes narrowed in focus. Whiskers, sensing his friend's excitement, followed suit, his tuxedo coat shimmering with curiosity.\n",
            "\n",
            "The squirrel, unaware of the danger lurking below, continued its carefree antics. Mittens pounced with lightning speed, but the squirrel was too quick. It leaped nimbly out of her reach, its bushy tail swishing behind it in amusement.\n",
            "\n",
            "Not to be\n",
            "________________________________________________________________________________\n",
            " deterred, Mittens gave chase, her tiny paws pattering along the grass. Whiskers, despite his more dignified nature, found himself caught up in the thrill of the hunt and joined in the pursuit.\n",
            "\n",
            "The chase led them through a rose garden, where the sweet fragrance of blooming flowers mingled with the sound of their laughter. They darted under trellises, weaving between fragrant petals and blooming buds.\n",
            "\n",
            "Finally, as the sun began to dip below the horizon, Mittens and Whiskers cornered the squirrel in a thicket of honeysuckle. The squirrel, its eyes wide with fear, scampered up a tree trunk and disappeared into a hole in the bark.\n",
            "\n",
            "Mittens and Whiskers abandoned their chase, their spirits still soaring. They had enjoyed their adventure, and the bond between them had grown even stronger. As they made their way back to the cottage, the setting sun cast a warm glow upon their furry faces, illuminating their playful hearts.\n",
            "________________________________________________________________________________\n",
            "==========not blocked!==========\n",
            "==========not blocked!==========\n"
          ]
        }
      ],
      "source": [
        "import asyncio\n",
        "\n",
        "async def get_response():\n",
        "    async for chunk in await model.generate_content_async(\"Write a cute story about cats.\", stream=True):\n",
        "      print(chunk.text)\n",
        "      print(\"_\"*80)\n",
        "\n",
        "async def something_else():\n",
        "    for i in range(5):\n",
        "        print(\"==========not blocked!==========\")\n",
        "        await asyncio.sleep(3)\n",
        "\n",
        "async def async_demo():\n",
        "    # Create tasks\n",
        "    task1 = asyncio.create_task(get_response())\n",
        "    task2 = asyncio.create_task(something_else())\n",
        "\n",
        "    # Wait for tasks to complete\n",
        "    await asyncio.gather(task1, task2)\n",
        "\n",
        "# Jupyter notebooks handle event loops for you, so await directly\n",
        "await async_demo()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Streaming.ipynb",
      "toc_visible": true
    },
    "google": {
      "image_path": "/site-assets/images/share.png",
      "keywords": [
        "examples",
        "googleai",
        "samplecode",
        "python",
        "embed",
        "function"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
